package com.yuba.common.ai.utils;

import com.yuba.common.ai.domain.VectorData;
import com.yuba.common.core.exception.ServiceException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;

/**
 * 向量数据库工具类型
 */
@Component
public class VectorStoreUtil {

    private static final Logger log = LoggerFactory.getLogger(VectorStoreUtil.class);
    private final VectorStore vectorStore;

    public VectorStoreUtil(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    public boolean save(List<VectorData> data) {
        // 拼装数据 按 category划分data内容
        HashMap<String, StringBuilder> map = new HashMap<>();
        for (VectorData vectorData : data) {
            String categoryLabel = vectorData.getCategoryLabel();
            String content = vectorData.getContent();

            if (!map.containsKey(categoryLabel)) {
                map.put(categoryLabel, new StringBuilder().append(categoryLabel).append("\r\n\t"));
            }

            map.get(categoryLabel).append(vectorData.getTitle()).append(content).append("\r\n");
        }

        // map中的转成String
        List<String> list = map.values().stream().map(StringBuilder::toString).toList();
        String vectorData = String.join("", list);

        try {
            ByteArrayResource resource = new ByteArrayResource(vectorData.getBytes());
            // 转成文档
            TikaDocumentReader reader = new TikaDocumentReader(resource);
            List<Document> documents = new TokenTextSplitter().apply(reader.read());

            // 存储向量数据库
            vectorStore.add(documents);
        } catch (Exception e) {
            log.error("同步异常信息",  e);
            throw new ServiceException("同步异常");
        }
        return true;
    }

}
